// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/expression/vector_function.h>
#include <pollux/testing/vector/vector_test_base.h>

namespace kumo::pollux::test {

/// Taking a flat or constant array vector, wraps its element vector in a
/// dictionary that reverses the order of rows.
class TestingDictionaryArrayElementsFunction : public exec::VectorFunction {
 public:
  TestingDictionaryArrayElementsFunction() {}

  static void register_function() {
    exec::registerVectorFunction(
        "testing_dictionary_array_elements",
        test::TestingDictionaryArrayElementsFunction::signatures(),
        std::make_unique<test::TestingDictionaryArrayElementsFunction>(),
        exec::VectorFunctionMetadataBuilder()
            .defaultNullBehavior(false)
            .build());
  }

  void apply(
      const SelectivityVector& /*rows*/,
      std::vector<VectorPtr>& args,
      const TypePtr& /*outputType*/,
      exec::EvalCtx& context,
      VectorPtr& result) const override {
    ArrayVector* array{nullptr};
    if (args[0]->encoding() == VectorEncoding::Simple::ARRAY) {
      array = args[0]->as<ArrayVector>();
    } else if (args[0]->is_constant_encoding()) {
      array = args[0]
                  ->as<ConstantVector<pollux::ComplexType>>()
                  ->value_vector()
                  ->as<ArrayVector>();
    }
    POLLUX_CHECK_NOT_NULL(
        array, "Input must be a flat or constant array vector.");

    auto elements = array->elements();
    const auto size = elements->size();
    auto indices = make_indices_in_reverse(size, context.pool());

    result = std::make_shared<ArrayVector>(
        context.pool(),
        array->type(),
        nullptr,
        array->size(),
        array->offsets(),
        array->sizes(),
        BaseVector::wrap_in_dictionary(nullptr, indices, size, elements),
        0);
  }

  static std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {
    // array(T) -> array(T)
    return {exec::FunctionSignatureBuilder()
                .typeVariable("T")
                .returnType("array(T)")
                .argumentType("array(T)")
                .build()};
  }
};

} // namespace kumo::pollux::test
